import sys
from pathlib import Path
PATH_ROOT = Path(__file__).parents[4]
sys.path.append(str(PATH_ROOT))
PATH_LOGS = PATH_ROOT / "manipulation/logs"
PATH_LOGS.mkdir(exist_ok=True)

import cv2
import numpy as np
from rclpy.node import Node
from manipulation.scripts.detection.cameras import BaseCamera
from sensor_msgs.msg import Image

import subprocess
import time

class DummyCamera(BaseCamera):
    """
    This is a dummy camera that subscribes to ROS2 topics to get RGB and Depth images.
    Which need rosbag play or real camera node to publish images.
    """
    dist_coefs = {  # D435 kuavo robot
        '1920x1080': [0, 0, 0, 0],
        '1280x720': [0, 0, 0, 0],
        '640x480': [0, 0, 0, 0]
    }
    K = np.array([  # D435 kuavo robot
        [607.8534 ,   0.     , 327.83737],
        [  0.     , 607.93445, 252.96724],
        [  0.     ,   0.     ,   1.     ]
    ], dtype=np.float32)
    dist = [0, 0, 0, 0]

    def __init__(self, rgb_width=640, rgb_height=480, fps=15, only_rgb=False, depth_width=640, depth_height=480, **kwargs):
        super().__init__(rgb_width, rgb_height, fps, only_rgb, depth_width, depth_height, **kwargs)
        self.node: Node = kwargs['node']
        self.node.create_subscription(Image, '/camera/color/image_raw', self.callback_rgb, 10)
        self.node.create_subscription(Image, '/camera/depth/image_raw', self.callback_depth, 10)
        self.last_update = [False, False]
        self.rgb_img = np.zeros((rgb_height, rgb_width, 3), np.uint8)
        self.depth_img = np.zeros((rgb_height, rgb_width), np.uint16)
    
    def callback_rgb(self, msg: Image):
        self.last_update[0] = True
        self.rgb_img = np.frombuffer(msg.data, dtype=np.uint8).reshape(msg.height, msg.width, -1)
    
    def callback_depth(self, msg: Image):
        self.last_update[1] = True
        self.depth_img = np.frombuffer(msg.data, dtype=np.uint16).reshape(msg.height, msg.width)
    
    def get_frame(self) -> tuple[np.ndarray, np.ndarray]:
        while True:
            if all(self.last_update):
                self.last_update = [False, False]
                break
            time.sleep(0.001)
        return self.rgb_img.copy(), self.depth_img.copy()
    
def debug_camera(width=640, height=480, fps=15, only_rgb=False, save=False):
    camera = DummyCamera(rgb_width=width, rgb_height=height, fps=fps, only_rgb=only_rgb)
    if save:
        path_video = PATH_LOGS / time.strftime("%Y%m%d_%H%M%S.avi")
        writer = cv2.VideoWriter(str(path_video), cv2.VideoWriter_fourcc(*'XVID'), fps=fps, frameSize=(width, height))
        print("path_video location:", path_video)
    while True:
        rgb_img, depth_img = camera.get_frame()
        depth_normalized = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX)
        depth_normalized = depth_normalized.astype(np.uint8)
        
        if save:
            writer.write(rgb_img)
        print(rgb_img.shape, depth_img.min(), depth_img.max())
        cv2.imshow("debug rgb", rgb_img)
        cv2.imshow("debug depth", depth_normalized)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    writer.release()
    
    
if __name__ == '__main__':
    debug_camera(save=False)
    # debug_camera(save=True)
